# _*_ coding:utf-8 _*_
import os
import subprocess
import time
import confman
import device_detect
try:
    import cv2
except:
    print("您尚未安装OpenCV，正在自动安装。此过程中可能需要您输入密码。")
    os.system("sudo pip3 install opencv-python")
    import cv2

ERR_CANNOT_OPEN_SRC = 1
ERR_BLANK_SRC = 2
ERR_NO_FFMPEG = 3
ERR_FFMPEG_ERROR = 4
ERR_TRANSMITION_FAILED = 5
ERR_CANNOT_READ_SRC = 6
ERR_MSG=["正常退出","无法打开视频源","未指定视频源","FFmpeg启动失败","FFmpeg错误","传输错误","无法读取视频帧"]

def log(content,lvl=0):
    if(lvl==0):
        print("[I][{}]{}".format(time.strftime('%H:%M:%S', time.localtime()),content))
    elif(lvl==1):
        print("[W][{}]{}".format(time.strftime('%H:%M:%S', time.localtime()),content))
    elif(lvl==2):
        print("[E][{}]{}".format(time.strftime('%H:%M:%S', time.localtime()),content))

def imgproc(img):
    img=cv2.flip(img,0)
    return img

# ffmpeg 推流
class FFmpegPusher(object):
    def __init__(self, rtmpaddr, videoid = 0, width = 1280, height = 720, 
                show_window = True, create_log = True):
        self.rtmpUrl = rtmpaddr
        self.video_stream_path = videoid
        self.WIDTH = width
        self.HEIGHT = height
        self.stat = True
        self.display = show_window
        self.genlog = create_log

    def open_opencv(self):
        cap = cv2.VideoCapture(self.video_stream_path)
        # 设置摄像头设备分辨率
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.WIDTH)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.HEIGHT)
        # 设置摄像头设备帧率,如不指定,默认30
        #cap.set(cv2.CAP_PROP_FPS, self.FPS)
        # 解决延迟
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
        return cap

    def open_ffmpeg(self):
        global log
        if(self.video_stream_path==None or str(self.video_stream_path).strip()==''):
            log("未指定视频源",2)
            return ERR_BLANK_SRC
        log("正在打开视频源{}……".format(self.video_stream_path))
        cap = self.open_opencv()
        if(not cap.isOpened()): 
            log("视频源{}无法打开！".format(self.video_stream_path),2)
            log("请检查文件是否存在，摄像头名称是否正确，设备连接是否可靠。",2)
            return ERR_CANNOT_OPEN_SRC
        else:
            log("视频源{}已经打开".format(self.video_stream_path))
        #fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # ffmpeg command
        command = ['ffmpeg','-y','-re',
        '-f', 'rawvideo',
        '-vcodec','rawvideo',
        '-pix_fmt', 'bgr24',
        '-s', "{}x{}".format(width, height),
        #'-r', str(fps),
        '-i', '-',
        '-c:v', 'libx264',
        '-pix_fmt', 'yuv420p',
        '-preset', 'ultrafast',
        '-f', 'flv', 
        self.rtmpUrl]

        log("将使用下列FFmpeg命令启动："+' '.join(command))
        
        # 管道配置
        # self.p = sp.Popen(command, stdin=sp.PIPE, shell=True)
        p = None
        if(self.genlog):
            tm = time.strftime('%Y%m%d_%H%M%S', time.localtime())
            mlog=open("pusher"+tm+".log",'w')
            try:
                p = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=mlog, stderr=mlog)
            except:
                return ERR_NO_FFMPEG
        else:
            try:
                p = subprocess.Popen(command, stdin=subprocess.PIPE)
            except:
                return ERR_NO_FFMPEG
        log("FFmpeg已启动，正在推流……")
        # read webcamera
        ret_code = 0
        while(cap.isOpened() and p.poll()==None):
            if(p.poll()!=None):
                log("FFmpeg异常退出！",2)
                ret_code = ERR_FFMPEG_ERROR
                break
            ret, frame = cap.read()
            frame=imgproc(frame)
            if not ret:
                log("摄像头错误。无法抓取帧！",2)
                ret_code = ERR_CANNOT_READ_SRC
                break
            elif not self.stat:
                p.kill()
                log("收到退出信号，视频上传已终止。")
                break
            try:
                p.stdin.write(frame.tobytes())
            except KeyboardInterrupt:
                log("收到用户的键盘退出指令，程序将结束。")
            except:
                p.kill()
                log("视频传输错误。请检查推流地址是否有效，网络连接是否可靠。",2)
                ret_code = ERR_TRANSMITION_FAILED
                break
            if(self.display): 
                cv2.imshow(self.rtmpUrl,frame)
                if(cv2.waitKey(1) & 0xFF==ord('q')): break
        
        if(self.display): cv2.destroyAllWindows()
        return ret_code

    # 关闭直播
    def close_ffmpeg(self):
        self.stat = False
        self.open_ffmpeg()

CFG_PATH = "pusher.ini"
DEF_CFG = '''[basic]
address = rtmp://1.13.152.174/live/test2            #视频录制地址，可以是rtmp服务器，也可以是文件
device  = 0                                         #设备名称，0表示默认设备
#特别提醒：在linux下，建议使用v4l2-ctl --list-devices查看设备名称
#对于Intel RealSense深度相机，请填入$RealSenseDepth$表示深度图像，填入$RealSenseRGB$表示普通图像
width   = 1280                                      #视频宽度
height  = 720                                       #视频高度
display = hide                                      #是否显示窗口，show为显示，hide为不显示
log     = true                                      #是否创建日志，true/false'''

def str2int(val:str):
    try:val=int(val)
    except:pass
    return val

def main():
    log("欢迎使用推流软件！请使用Ctrl+C退出。")
    opt=confman.read_or_create(CFG_PATH,DEF_CFG)['basic']
    if(opt==None):return
    try:
        addr,dev,width,height=opt["address"],opt["device"],int(opt["width"]),int(opt["height"])
        dev=str2int(dev)
        if(dev=="$RealSenseRGB$"):dev=device_detect.find_realsense_camera("RGB")
        if(dev=="$RealSenseDepth$"):dev=device_detect.find_realsense_camera("Depth")
        log("============推流信息============")
        log("设备：{}".format(dev))
        log("画面大小：{}x{}".format(width,height))
        log("地址：{}".format(addr))
        log("================================")
        
    except:
        log("读取配置时发生错误，程序将结束。",2)
        return
    try:
        fr = FFmpegPusher(addr,dev,width,height,
                        opt["display"].lower()=="true",
                        opt["log"].lower()=="true",)
        ret_code=-1
        while(ret_code==-1 or ret_code>=3):
            ret_code=fr.open_ffmpeg()
            if(ret_code==0):
                log("视频推流结束")
                break
            else:
                log("视频推流异常结束：{}".format(ERR_MSG[ret_code]),2)
                if(ret_code==ERR_NO_FFMPEG):
                    log("软件未检测到FFmpeg。正在安装FFmpeg，您可能需要输入密码。")
                    os.system("sudo apt-get install ffmpeg")
                elif(ret_code>=3):
                    log("软件将在10s后尝试重连。")
                    log("您可以使用Ctrl+C终止这一过程。")
                    time.sleep(10)
    except KeyboardInterrupt:
        log("收到用户的键盘退出指令，程序将结束。")
    except:
        log("无法创建视频上传实例，程序将结束。",2)

if __name__ == "__main__":
    main()
